package com.wunong.smart.bigdata.flink.job;

import com.wunong.smart.bigdata.flink.data.TimestampData;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.CoGroupFunction;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.util.Collector;

import java.time.Duration;

/**
 * coGroup calculator
 *
 * @author created by zealot.zt
 */
public class CoGroupJob {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment()
                .setParallelism(1);

        DataStream<TimestampData> stream1 = env.fromElements(
                TimestampData.create(1),
                TimestampData.create(2),
                TimestampData.create(3));

        long mills = System.currentTimeMillis();
        DataStream<TimestampData> stream2 = env.fromElements(
                TimestampData.create(1, 10),
                TimestampData.create(2, 20, mills - 60 * 1000L),
                TimestampData.create(3, 30, mills - 2 * 60 * 1000L),
                TimestampData.create(4, 20)
        );

        // 时间水印
        WatermarkStrategy<TimestampData> watermarkStrategy = getWatermarkStrategy();

        // 分组合并
        SingleOutputStreamOperator<TimestampData> keyStream2 = stream2.assignTimestampsAndWatermarks(watermarkStrategy);
        stream1.assignTimestampsAndWatermarks(watermarkStrategy)
                .coGroup(keyStream2)
                .where(new KeySelector<TimestampData, String>() {
                    @Override
                    public String getKey(TimestampData value) throws Exception {
                        return String.valueOf(value.getId());
                    }
                })
                .equalTo(new KeySelector<TimestampData, String>() {
                    @Override
                    public String getKey(TimestampData value) throws Exception {
                        return String.valueOf(value.getId());
                    }
                })
                // 时间窗口，需要包含当前基准时间、水印时间
                .window(TumblingEventTimeWindows.of(Time.minutes(5)))
                .apply(new CoGroupFunction<TimestampData, TimestampData, String>() {
                    @Override
                    public void coGroup(Iterable<TimestampData> first, Iterable<TimestampData> second,
                                        Collector<String> out) throws Exception {
                        TimestampData firstMaxData = null;

                        // 选最大数据
                        for (TimestampData f : first) {
                            if (firstMaxData == null || firstMaxData.getTimestamp() < f.getTimestamp()) {
                                firstMaxData = f;
                            }
                        }

                        // 选最大数据
                        TimestampData secondMaxData = null;
                        for (TimestampData s : second) {
                            if (secondMaxData == null || secondMaxData.getTimestamp() < s.getTimestamp()) {
                                secondMaxData = s;
                            }
                        }

                        StringBuilder builder = new StringBuilder(64);
                        if (firstMaxData != null) {
                            builder.append(firstMaxData.getId()).append(":").append(firstMaxData.getData()).append(";");
                        }

                        if (secondMaxData != null) {
                            builder.append(secondMaxData.getId()).append(":").append(secondMaxData.getData());
                        }

                        out.collect(builder.toString());
                    }
                }).print("Result");

        env.execute(CoGroupJob.class.getName());
    }

    /**
     * 时间水印
     *
     * @return
     */
    public static WatermarkStrategy<TimestampData> getWatermarkStrategy() {
        // 时间水印，对事件时间内保证有序
        final Duration widowDuration = Duration.ofSeconds(300 * 3 / 5);
        WatermarkStrategy<TimestampData> watermarkStrategy = WatermarkStrategy
                .<TimestampData>forBoundedOutOfOrderness(widowDuration)
                .withTimestampAssigner(((message, timestamp) -> message.getTimestamp()));
        return watermarkStrategy;
    }

}